# 在 lesson 02 代码的基础上 找出既满足在兴趣区域内 又满足色彩阈值的进行标注

# 导入所需包
from functools import reduce
import matplotlib.pyplot as plt
import matplotlib.image as mping
import numpy as np

# 传入图片
img = mping.imread('test2.jpg')


# 输出图片格式与大小
print('This image is: ', type(img), 'with dimensions: ', img.shape)

# 复制图像的宽高
y = img.shape[0]
x = img.shape[1]

# 复制颜色选择
color_select = np.copy(img)
line_image = np.copy(img)

# 定义色彩阈值
red_threshold = 200
green_threshold = 200
blue_threshold = 200

rgb_threshold = [red_threshold, green_threshold, blue_threshold]

# print(color_select)

# 定一个三角形的区域 
# 默认原点 (0, 0) 在左上角
left_bottom = [0, 700]
right_bottom = [1200, 700]
apex = [600, 320]

# 拟合线 (y = Ax + B) 识别ROI的三边区域
# np.polyfit() 用来返回系数 A B
fit_left = np.polyfit((left_bottom[0], apex[0]), 
                        (left_bottom[1], apex[1]), 1)

fit_right = np.polyfit((right_bottom[0], apex[0]), 
                        (right_bottom[1], apex[1]), 1)

fit_bottom = np.polyfit((left_bottom[0], right_bottom[0]), 
                        (left_bottom[1], right_bottom[1]), 1)


color_thresholds = (img[:,:,0] < rgb_threshold[0]) | \
                    (img[:,:,1] < rgb_threshold[1]) | \
                    (img[:,:,2] < rgb_threshold[2])


XX, YY = np.meshgrid(np.arange(0, x), np.arange(0, y))
region_thresholds = (YY > (XX*fit_left[0] + fit_left[1])) & \
                    (YY > (XX*fit_right[0] + fit_right[1])) & \
                    (YY < (XX*fit_bottom[0] + fit_bottom[1]))

# 不满足为黑色
color_select[color_thresholds | ~region_thresholds] = [0, 0, 0]

# 颜色和区域都满足 定位红色
line_image[~color_thresholds & region_thresholds] = [255, 0, 0]

# 显示图像并显示区域和颜色选择
plt.imshow(img)
x = [left_bottom[0], right_bottom[0], apex[0], left_bottom[0]]
y = [left_bottom[1], right_bottom[1], apex[1], left_bottom[1]]
plt.plot(x, y, 'b--', lw=4)

plt.imshow(color_select)
plt.imshow(line_image)
plt.show()

